package com.bw.sparksql1.job2

import org.apache.log4j.{Level, Logger}
import org.apache.spark.sql.expressions.{MutableAggregationBuffer, UserDefinedAggregateFunction}
import org.apache.spark.sql.types._
import org.apache.spark.sql.{DataFrame, Row, SparkSession}

/**
  *
  *  自定义UDAF：自定聚合函数
  *  需求：求平均薪水
  *
  */
object Job13 extends UserDefinedAggregateFunction {

  /**
    * 定义输入的数据类型
    * @return
    */
  override def inputSchema: StructType = StructType(
    StructField("salary",DoubleType,true) :: Nil
  )

  /**
    * 定义输出的数据类型
    * @return
    */
  override def dataType: DataType = StructType(
    StructField("price_avg",DoubleType,true) :: Nil
  )


  /**
    * 定义辅助字段：
    * 辅助字段一：total表示总金额
    * 辅助字段二：count表示总人数
    */
  override def bufferSchema: StructType = StructType(
    StructField("total",DoubleType,true) ::
      StructField("count",IntegerType,true) ::
    Nil
  )

  /**
    * 最后的目标计算函数
    */
  override def evaluate(buffer: Row): Any = {
    val total = buffer.getDouble(0)
    val count = buffer.getInt(1)
    total / count
  }

  /**
    *
    * 初始化辅助字段
    */
  override def initialize(buffer: MutableAggregationBuffer): Unit = {
    buffer.update(0,0.0)
    buffer.update(1,0)
  }

  /**
    * 跟新辅助字段的值  -> 局部的，表示在一个分区中的
    * @param buffer
    * @param input
    */
  override def update(buffer: MutableAggregationBuffer, input: Row): Unit = {
    //获取上一次的值
    val lastTotal = buffer.getDouble(0)
    val lastCount = buffer.getInt(1)

    //获取当前的salary
    val currentSalary = input.getDouble(0)

    //更新
    buffer.update(0,lastTotal+currentSalary)
    buffer.update(1,lastCount+1)
  }

  /**
    * 全局的 ->
    *  partition1 和 partiton2 需要合并
    */
  override def merge(buffer1: MutableAggregationBuffer, buffer2: Row): Unit = {
    //当前分区
    val total1 = buffer1.getDouble(0)
    val count1 = buffer1.getInt(1)

    //其他分区
    val total2 = buffer2.getDouble(0)
    val count2 = buffer2.getInt(1)

    buffer1.update(0,total1+total2)
    buffer1.update(1,count1+count2)



  }

  //输入和输出的数据类型是否一致
  override def deterministic: Boolean = true




}